package com.sugon.xuanyuan;

import ai.djl.util.JsonUtils;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.ocr.config.PlateDetModelConfig;
import cn.smartjavaai.ocr.config.PlateRecModelConfig;
import cn.smartjavaai.ocr.entity.PlateInfo;
import cn.smartjavaai.ocr.enums.PlateDetModelEnum;
import cn.smartjavaai.ocr.enums.PlateRecModelEnum;
import cn.smartjavaai.ocr.factory.PlateModelFactory;
import cn.smartjavaai.ocr.model.plate.PlateDetModel;
import cn.smartjavaai.ocr.model.plate.PlateRecModel;
import lombok.extern.slf4j.Slf4j;
import org.apache.spark.sql.api.java.UDF1;

import java.util.List;

@Slf4j
public class LPRUDF implements UDF1<byte[], String> {

    public static final String basePath = BaseConstants.basePath;

    public static final DeviceEnum device = BaseConstants.device;

    public static PlateRecModel recModel;

    static {
        log.info("init lpr model---------------------------------------------------");
        recModel=getPlateRecModel() ;
        log.info("init lpr model---------------------------------------------------");
    }

    public static PlateDetModel getPlateDetModel() {
        PlateDetModelConfig config = new PlateDetModelConfig();
        config.setModelEnum(PlateDetModelEnum.YOLOV5);
        config.setModelPath(basePath+"/lprmodels/lpr_det/yolov5_plate_detect.onnx");
        config.setDevice(device);
        return PlateModelFactory.getInstance().getDetModel(config);
    }

    /**
     * 获取车牌识别模型
     * @return
     */

    public static PlateRecModel getPlateRecModel() {
        PlateRecModelConfig recModelConfig = new PlateRecModelConfig();
        recModelConfig.setModelEnum(PlateRecModelEnum.PLATE_REC_CRNN);
        recModelConfig.setModelPath(basePath+"/lprmodels/lpr_rec/plate_rec_color.onnx");
        recModelConfig.setPlateDetModel(getPlateDetModel());
        return PlateModelFactory.getInstance().getRecModel(recModelConfig);
    }

    @Override
    public String call(byte[] bytes) throws Exception {
        try {
            R<List<PlateInfo>> result = recModel.recognize(bytes);
            if(result.isSuccess()){

                return JsonUtils.toJson(result.getData());
            }else{
                throw  new RuntimeException(result.getMessage());
            }
        }catch (Exception e){
            log.error("识别车牌失败",e);
            throw new RuntimeException(e);
        }
    }
}
